import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class StockRevenuePrediction:
    def __init__(self):
        self.file_path = '上市公司营收数据.xlsx'
        self.df = None
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
        self.scaler = StandardScaler()
        self.models = {}
        self.predictions = {}
        self.metrics = {}
    
    def load_data(self):
        """加载和初步分析数据"""
        print("\n===== 1. 数据加载与初步分析 =====")
        self.df = pd.read_excel(self.file_path)
        print(f"数据加载完成，共有 {self.df.shape[0]} 行，{self.df.shape[1]} 列")
        print("\n数据列名：")
        for i, col in enumerate(self.df.columns, 1):
            print(f"{i}. {col}")
        
        # 查看数据基本统计信息
        print("\n数据基本统计描述：")
        print(self.df.describe())
        
        # 检查缺失值
        print("\n缺失值统计：")
        missing_data = self.df.isnull().sum()
        if missing_data.sum() == 0:
            print("✓ 数据无缺失值")
        else:
            print(missing_data[missing_data > 0])
    
    def data_preprocessing(self):
        """数据预处理和特征工程"""
        print("\n===== 2. 数据预处理和特征工程 =====")
        
        # 选择特征变量（排除证券代码和证券简称）
        feature_cols = [col for col in self.df.columns if col not in ['证券代码', '证券简称', '营业收入']]
        target_col = '营业收入'
        
        X = self.df[feature_cols]
        y = self.df[target_col]
        
        # 数据分割
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            X, y, test_size=0.2, random_state=42)
        
        print(f"训练集大小：{self.X_train.shape[0]}，测试集大小：{self.X_test.shape[0]}")
        
        # 特征标准化
        self.X_train_scaled = self.scaler.fit_transform(self.X_train)
        self.X_test_scaled = self.scaler.transform(self.X_test)
        
        # 创建特征名称映射
        self.feature_names = feature_cols
        
        print("✓ 数据预处理完成")
    
    def build_models(self):
        """构建多种预测模型"""
        print("\n===== 3. 构建预测模型 =====")
        
        # 线性回归
        print("\n3.1 训练线性回归模型...")
        lr_model = LinearRegression()
        lr_model.fit(self.X_train_scaled, self.y_train)
        self.models['线性回归'] = lr_model
        
        # Ridge回归
        print("3.2 训练Ridge回归模型...")
        ridge_model = Ridge(alpha=1.0)
        ridge_model.fit(self.X_train_scaled, self.y_train)
        self.models['Ridge回归'] = ridge_model
        
        # Lasso回归
        print("3.3 训练Lasso回归模型...")
        lasso_model = Lasso(alpha=0.1)
        lasso_model.fit(self.X_train_scaled, self.y_train)
        self.models['Lasso回归'] = lasso_model
        
        # 随机森林回归
        print("3.4 训练随机森林回归模型...")
        rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
        rf_model.fit(self.X_train, self.y_train)  # 随机森林不需要标准化
        self.models['随机森林'] = rf_model
        
        # 梯度提升回归
        print("3.5 训练梯度提升回归模型...")
        gb_model = GradientBoostingRegressor(n_estimators=100, random_state=42)
        gb_model.fit(self.X_train, self.y_train)
        self.models['梯度提升'] = gb_model
        
        print("✓ 所有模型训练完成")
    
    def evaluate_models(self):
        """评估所有模型的性能"""
        print("\n===== 4. 模型评估 =====")
        
        print("\n各模型性能指标：")
        print("-" * 80)
        print(f"{'模型名称':<12} {'MSE':<15} {'RMSE':<15} {'MAE':<15} {'R²':<10}")
        print("-" * 80)
        
        for name, model in self.models.items():
            if name in ['线性回归', 'Ridge回归', 'Lasso回归']:
                y_pred = model.predict(self.X_test_scaled)
            else:
                y_pred = model.predict(self.X_test)
            
            self.predictions[name] = y_pred
            
            mse = mean_squared_error(self.y_test, y_pred)
            rmse = np.sqrt(mse)
            mae = mean_absolute_error(self.y_test, y_pred)
            r2 = r2_score(self.y_test, y_pred)
            
            self.metrics[name] = {'MSE': mse, 'RMSE': rmse, 'MAE': mae, 'R²': r2}
            
            print(f"{name:<12} {mse:<15.2e} {rmse:<15.2e} {mae:<15.2e} {r2:<10.4f}")
        
        print("-" * 80)
        
        # 找出最优模型
        best_model_name = max(self.metrics, key=lambda x: self.metrics[x]['R²'])
        self.best_model = self.models[best_model_name]
        self.best_model_name = best_model_name
        
        print(f"\n最优模型：{best_model_name}")
        print(f"最优模型 R² 得分：{self.metrics[best_model_name]['R²']:.4f}")
    
    def feature_importance_analysis(self):
        """分析特征重要性"""
        print("\n===== 5. 特征重要性分析 =====")
        
        # 对于不同类型的模型，使用不同的特征重要性计算方法
        if self.best_model_name in ['随机森林', '梯度提升']:
            importances = self.best_model.feature_importances_
        else:  # 线性模型
            importances = np.abs(self.best_model.coef_)
        
        # 创建特征重要性DataFrame
        feature_importance = pd.DataFrame({
            '特征': self.feature_names,
            '重要性': importances
        })
        
        # 按重要性排序
        feature_importance = feature_importance.sort_values('重要性', ascending=False)
        
        print(f"\n{self.best_model_name} 模型的特征重要性排序：")
        for i, (idx, row) in enumerate(feature_importance.iterrows(), 1):
            print(f"{i}. {row['特征']}: {row['重要性']:.4f}")
        
        # 保存关键特征
        self.key_features = feature_importance.head(5)['特征'].tolist()
        print(f"\n影响上市公司营收的关键因素（前5个）：")
        for i, feature in enumerate(self.key_features, 1):
            print(f"{i}. {feature}")
    
    def visualize_results(self):
        """可视化结果"""
        print("\n===== 6. 结果可视化 =====")
        
        # 1. 各模型性能对比
        plt.figure(figsize=(12, 6))
        
        # R²对比
        plt.subplot(1, 2, 1)
        models = list(self.metrics.keys())
        r2_scores = [self.metrics[m]['R²'] for m in models]
        plt.bar(models, r2_scores, color='skyblue')
        plt.title('各模型 R² 得分对比')
        plt.ylabel('R²')
        plt.ylim(0, 1)
        plt.xticks(rotation=45)
        
        # RMSE对比
        plt.subplot(1, 2, 2)
        rmse_scores = [self.metrics[m]['RMSE'] for m in models]
        plt.bar(models, rmse_scores, color='lightgreen')
        plt.title('各模型 RMSE 对比')
        plt.ylabel('RMSE')
        plt.xticks(rotation=45)
        
        plt.tight_layout()
        plt.savefig('模型性能对比.png', dpi=300, bbox_inches='tight')
        
        # 2. 最优模型的特征重要性图
        plt.figure(figsize=(10, 6))
        if self.best_model_name in ['随机森林', '梯度提升']:
            importances = self.best_model.feature_importances_
        else:
            importances = np.abs(self.best_model.coef_)
        
        indices = np.argsort(importances)[::-1]
        plt.title(f'{self.best_model_name} 特征重要性')
        plt.bar(range(len(indices)), importances[indices])
        plt.xticks(range(len(indices)), [self.feature_names[i] for i in indices], rotation=90)
        plt.tight_layout()
        plt.savefig('特征重要性分析.png', dpi=300, bbox_inches='tight')
        
        # 3. 预测值与实际值对比图
        plt.figure(figsize=(8, 6))
        best_pred = self.predictions[self.best_model_name]
        plt.scatter(self.y_test, best_pred, alpha=0.5)
        plt.plot([self.y_test.min(), self.y_test.max()], [self.y_test.min(), self.y_test.max()], 'r--')
        plt.xlabel('实际营收')
        plt.ylabel('预测营收')
        plt.title(f'{self.best_model_name} 预测值与实际值对比')
        plt.tight_layout()
        plt.savefig('预测vs实际.png', dpi=300, bbox_inches='tight')
        
        print("✓ 可视化结果已保存")
    
    def generate_report(self):
        """生成实战报告"""
        print("\n===== 7. 生成实战报告 =====")
        
        report = """
# 上市公司营收预测实战报告

## 一、项目背景与目标

在金融领域，上市公司营收预测是一项重要的分析任务，它直接反映了公司的运营情况和未来发展趋势。
本项目旨在建立上市公司营收预测模型，评价模型性能，并探索影响上市公司营收的关键因素。

## 二、数据概览

### 2.1 数据基本信息
- 数据总量：{}行，{}列
- 数据完整性：无缺失值
- 时间范围：根据数据推测

### 2.2 数据字段说明

| 序号 | 字段名称 | 数据类型 | 说明 |
|-----|---------|---------|------|
""".format(self.df.shape[0], self.df.shape[1])
        
        # 添加字段说明
        for i, (col, dtype) in enumerate(zip(self.df.columns, self.df.dtypes), 1):
            report += f"| {i} | {col} | {dtype} | - |\n"
        
        # 数据统计概览
        report += """

### 2.3 数据统计概览

主要财务指标的统计信息：

| 指标 | 平均值 | 标准差 | 最小值 | 最大值 |
|-----|-------|-------|-------|-------|
"""
        
        # 添加关键指标统计
        key_metrics = ['营业收入', '营业成本', '资产总计']
        for metric in key_metrics:
            if metric in self.df.columns:
                mean_val = self.df[metric].mean()
                std_val = self.df[metric].std()
                min_val = self.df[metric].min()
                max_val = self.df[metric].max()
                report += f"| {metric} | {mean_val:.2e} | {std_val:.2e} | {min_val:.2e} | {max_val:.2e} |\n"
        
        # 三、模型构建与评估
        report += """

## 三、模型构建与评估

### 3.1 模型选择

本项目构建了以下五种预测模型：
1. 线性回归（Linear Regression）
2. Ridge回归（岭回归）
3. Lasso回归（套索回归）
4. 随机森林回归（Random Forest）
5. 梯度提升回归（Gradient Boosting）

### 3.2 模型性能评估

各模型性能指标对比：

| 模型名称 | MSE | RMSE | MAE | R² 得分 |
|---------|-----|------|-----|--------|
"""
        
        for name, metrics in self.metrics.items():
            report += f"| {name} | {metrics['MSE']:.2e} | {metrics['RMSE']:.2e} | {metrics['MAE']:.2e} | {metrics['R²']:.4f} |\n"
        
        # 最优模型
        report += f"""

### 3.3 最优模型选择

根据R²得分，**{self.best_model_name}** 模型表现最佳，其R²得分为 **{self.metrics[self.best_model_name]['R²']:.4f}**。

## 四、特征重要性分析

### 4.1 关键影响因素

通过对最优模型的特征重要性分析，我们识别出以下影响上市公司营收的关键因素（按重要性排序）：
"""
        
        # 添加关键特征
        for i, feature in enumerate(self.key_features, 1):
            report += f"{i}. **{feature}**\n"
        
        # 五、结论与建议
        report += """

## 五、结论与建议

### 5.1 主要结论

1. **模型效果**：{self.best_model_name}在预测上市公司营收方面表现最佳，说明营收预测问题可能具有一定的非线性特征。

2. **关键因素**：从特征重要性分析可以看出，{self.key_features[0]}和{self.key_features[1]}等因素对营收有显著影响，这些因素反映了公司的整体规模和运营效率。

3. **数据质量**：数据集质量良好，无缺失值，为模型构建提供了可靠的基础。

### 5.2 业务建议

1. **投资决策**：投资者在评估公司价值时，可以重点关注上述关键财务指标，特别是{self.key_features[0]}指标。

2. **模型改进**：未来可以考虑引入更多的特征变量，如行业分类、宏观经济指标等，进一步提高预测精度。

3. **时间序列分析**：本项目采用的是截面数据，可以考虑收集时间序列数据，建立动态预测模型。

### 5.3 局限性

1. 本模型基于历史财务数据构建，无法完全捕捉市场突变和外部冲击的影响。
2. 不同行业的公司可能存在显著差异，行业特异性分析有待进一步深入。

## 六、附录

### 6.1 可视化结果

- **模型性能对比图**：直观展示了各模型的R²和RMSE指标
- **特征重要性分布图**：展示了各特征对预测结果的贡献程度
- **预测值与实际值散点图**：评估模型预测的准确性和偏差
        """
        
        # 保存报告
        with open('上市公司营收预测实战报告.md', 'w', encoding='utf-8') as f:
            f.write(report)
        
        print("✓ 实战报告已生成")
        print("报告文件：上市公司营收预测实战报告.md")
    
    def run(self):
        """运行完整的分析流程"""
        print("===== 上市公司营收预测实战演练 =====")
        self.load_data()
        self.data_preprocessing()
        self.build_models()
        self.evaluate_models()
        self.feature_importance_analysis()
        self.visualize_results()
        self.generate_report()
        print("\n===== 实战演练完成 =====")

if __name__ == "__main__":
    # 初始化并运行分析
    predictor = StockRevenuePrediction()
    predictor.run()